import cv2
import numpy as np
import pyautogui
import pydirectinput
from ultralytics import YOLO
import time
import random
import heapq

# 定义移动方向
DIRECTIONS = [(0, -1), (1, 0), (0, 1), (-1, 0)]  # 上，右，下，左
current_direction_index = 0

# 红色方块大小（步长）
STEP_SIZE = 16  # 假设红色方块是16x16像素

# 安全区域扩展，避免紧贴边缘
SAFETY_MARGIN = 1


def capture_screen(x, y, width, height):
    """
    捕获指定区域的屏幕截图
    :param x: 区域左上角 x 坐标
    :param y: 区域左上角 y 坐标
    :param width: 区域宽度
    :param height: 区域高度
    :return: 转换为 OpenCV 格式的截图
    """
    screenshot = pyautogui.screenshot(region=(x, y, width, height))
    return cv2.cvtColor(np.array(screenshot), cv2.COLOR_RGB2BGR)


def detect_objects(image, model):
    """
    使用 YOLOv8 模型对图像进行目标检测
    :param image: 待检测的图像
    :param model: 加载好的 YOLO 模型
    :return: 检测结果
    """
    return model(image)


def print_detection_results(results):
    """
    打印检测结果，并返回矿石和岩浆的位置列表
    :param results: 检测结果
    :return: 矿石位置列表和岩浆位置列表
    """
    detection_info = {
        '岩浆': [],
        '石块': [],
        '自己': []
    }
    for r in results:
        boxes = r.boxes
        for box in boxes:
            class_id = r.names[int(box.cls)]
            confidence = float(box.conf)
            x1, y1, x2, y2 = box.xyxy[0].cpu().numpy().astype(int)
            # print(f"检测到对象: {class_id}, 置信度: {confidence:.2f}, 位置: ({x1}, {y1}, {x2}, {y2})")
            if class_id in detection_info:
                detection_info[class_id].append((x1, y1, x2, y2))
    return detection_info['石块'], detection_info['岩浆'], detection_info['自己']


def get_center(x1, y1, x2, y2):
    """
    计算目标的中心位置，并确定以中心为基准的 10x10 像素方块的位置
    :param x1: 目标左上角 x 坐标
    :param y1: 目标左上角 y 坐标
    :param x2: 目标右下角 x 坐标
    :param y2: 目标右下角 y 坐标
    :return: 10x10 像素方块的左上角和右下角坐标
    """
    center_x = int((x1 + x2) / 2)
    center_y = int((y1 + y2) / 2)
    crop_x1 = max(center_x - 5, 0)
    crop_y1 = max(center_y - 5, 87)
    crop_x2 = min(center_x + 5, 842)  # 根据截图宽度调整
    crop_y2 = min(center_y + 5, 658)  # 根据截图高度调整
    print(f"人物方块位置: 左上角({crop_x1}, {crop_y1}) 右下角({crop_x2}, {crop_y2})")
    return crop_x1, crop_y1, crop_x2, crop_y2


def visualize_areas(image, ore_locations, magma_locations, x, y, width, height):
    # 创建颜色掩码
    color_mask = np.zeros_like(image)

    # 向下移动危险区域 100 像素
    moved_ore_locations = [(x1, y1 + 80, x2, y2 + 80) for x1, y1, x2, y2 in ore_locations]
    moved_magma_locations = [(x1, y1 + 100, x2, y2 + 100) for x1, y1, x2, y2 in magma_locations]

    # 标记危险区域（石块和岩浆）为红色
    for loc in moved_ore_locations + moved_magma_locations:
        x1, y1, x2, y2 = loc
        if y1 < height and y2 > 0:  # 确保部分区域在截图内
            y1 = max(y1, 0)
            y2 = min(y2, height)
            color_mask[y1 - y:y2 - y, x1 - x:x2 - x] = [0, 0, 255]  # 红色

    # 创建安全区域掩码
    safe_mask = np.ones((height, width), dtype=np.uint8) * 255
    for loc in moved_ore_locations + moved_magma_locations:
        x1, y1, x2, y2 = loc
        if y1 < height and y2 > 0:
            y1 = max(y1, 0)
            y2 = min(y2, height)
            safe_mask[y1 - y:y2 - y, x1 - x:x2 - x] = 0

    # 扩展危险区域，增加安全边距
    kernel = np.ones((SAFETY_MARGIN * 2 + 1, SAFETY_MARGIN * 2 + 1), np.uint8)
    safe_mask = cv2.erode(safe_mask, kernel, iterations=1)

    # 标记安全区域为绿色
    color_mask[safe_mask > 0] = [0, 255, 0]  # 绿色

    # 将颜色掩码与原图合并
    alpha = 0.4  # 透明度
    result = cv2.addWeighted(image, 1, color_mask, alpha, 0)

    return result, safe_mask


def create_grid_from_mask(safe_mask):
    """将安全掩码转换为网格表示"""
    height, width = safe_mask.shape
    grid_width = width // STEP_SIZE
    grid_height = height // STEP_SIZE

    grid = np.ones((grid_height, grid_width), dtype=np.uint8) * 255

    for y in range(grid_height):
        for x in range(grid_width):
            # 检查网格单元的中心点是否在安全区域内
            cell_x = x * STEP_SIZE + STEP_SIZE // 2
            cell_y = y * STEP_SIZE + STEP_SIZE // 2

            if cell_x < width and cell_y < height and safe_mask[cell_y, cell_x] == 0:
                grid[y, x] = 0

    return grid


def find_safe_edge(grid, current_pos):
    """
    在网格中寻找安全区域的边缘位置
    :param grid: 安全区域网格
    :param current_pos: 当前位置 (x, y) 网格坐标
    :return: 边缘位置坐标 (x, y) 网格坐标
    """
    height, width = grid.shape
    current_x, current_y = current_pos

    # 计算到各边缘的距离
    distances = {
        'left': current_x,
        'right': width - 1 - current_x,
        'top': current_y,
        'bottom': height - 1 - current_y
    }

    # 按距离排序
    sorted_edges = sorted(distances.items(), key=lambda x: x[1], reverse=True)

    # 检查每个边缘方向，选择最近的可达边缘
    for edge, _ in sorted_edges:
        if edge == 'left':
            for x in range(current_x, -1, -1):
                if grid[current_y, x] == 0:
                    if x + 1 < width:
                        return (x + 1, current_y)
                    else:
                        break
            return (0, current_y)

        elif edge == 'right':
            for x in range(current_x, width):
                if grid[current_y, x] == 0:
                    if x - 1 >= 0:
                        return (x - 1, current_y)
                    else:
                        break
            return (width - 1, current_y)

        elif edge == 'top':
            for y in range(current_y, -1, -1):
                if grid[y, current_x] == 0:
                    if y + 1 < height:
                        return (current_x, y + 1)
                    else:
                        break
            return (current_x, 0)

        elif edge == 'bottom':
            for y in range(current_y, height):
                if grid[y, current_x] == 0:
                    if y - 1 >= 0:
                        return (current_x, y - 1)
                    else:
                        break
            return (current_x, height - 1)

    # 如果没有找到边缘，返回当前位置
    return current_pos


def heuristic(a, b):
    """A*算法的启发函数（曼哈顿距离）"""
    return abs(a[0] - b[0]) + abs(a[1] - b[1])


def a_star_search(grid, start, goal):
    """
    使用A*算法在网格中寻找从起点到终点的路径
    :param grid: 安全区域网格
    :param start: 起点坐标 (x, y) 网格坐标
    :param goal: 终点坐标 (x, y) 网格坐标
    :return: 路径列表 (x, y) 网格坐标
    """
    start = (int(start[0]), int(start[1]))
    goal = (int(goal[0]), int(goal[1]))

    # 检查起点和终点是否在网格内且安全
    if start[0] < 0 or start[1] < 0 or start[0] >= grid.shape[1] or start[1] >= grid.shape[0] or grid[
        start[1], start[0]] == 0:
        print(f"起点 {start} 不在安全区域内")
        return []

    if goal[0] < 0 or goal[1] < 0 or goal[0] >= grid.shape[1] or goal[1] >= grid.shape[0] or grid[
        goal[1], goal[0]] == 0:
        print(f"终点 {goal} 不在安全区域内")
        return []

    neighbors = [(0, 1), (0, -1), (1, 0), (-1, 0)]  # 上下左右移动
    close_set = set()
    came_from = {}
    gscore = {start: 0}
    fscore = {start: heuristic(start, goal)}
    oheap = []

    heapq.heappush(oheap, (fscore[start], start))

    while oheap:
        current = heapq.heappop(oheap)[1]

        if current == goal:
            data = []
            while current in came_from:
                data.append(current)
                current = came_from[current]
            return [(pos[0], pos[1]) for pos in reversed(data)]

        close_set.add(current)

        for i, j in neighbors:
            neighbor = current[0] + i, current[1] + j

            # 检查是否超出边界
            if neighbor[0] < 0 or neighbor[1] < 0 or neighbor[0] >= grid.shape[1] or neighbor[1] >= grid.shape[0]:
                continue

            # 检查是否是障碍物
            if grid[neighbor[1], neighbor[0]] == 0:
                continue

            tentative_g_score = gscore[current] + 1  # 网格移动成本为1

            if neighbor in close_set and tentative_g_score >= gscore.get(neighbor, 0):
                continue

            if tentative_g_score < gscore.get(neighbor, 0) or neighbor not in [i[1] for i in oheap]:
                came_from[neighbor] = current
                gscore[neighbor] = tentative_g_score
                fscore[neighbor] = tentative_g_score + heuristic(neighbor, goal)
                heapq.heappush(oheap, (fscore[neighbor], neighbor))

    return []


def grid_to_pixel(grid_pos):
    """将网格坐标转换为像素坐标"""
    return (grid_pos[0] * STEP_SIZE + STEP_SIZE // 2,
            grid_pos[1] * STEP_SIZE + STEP_SIZE // 2)


def pixel_to_grid(pixel_pos, width, height):
    """将像素坐标转换为网格坐标"""
    grid_x = min(pixel_pos[0] // STEP_SIZE, (width // STEP_SIZE) - 1)
    grid_y = min(pixel_pos[1] // STEP_SIZE, (height // STEP_SIZE) - 1)
    return (grid_x, grid_y)


def move_in_safe_area(x, y, width, height, model):
    # 定义全局变量以跟踪当前方向索引和路径
    global current_direction_index

    # 捕获屏幕指定区域的截图
    screenshot = capture_screen(x, y, width, height)

    # 使用模型检测截图中的对象
    results = detect_objects(screenshot, model)

    # 打印检测结果并获取矿石、岩浆和角色的位置
    ore_locations, magma_locations, role_locations = print_detection_results(results)

    # 可视化检测区域并标记安全区域
    marked_image, safe_mask = visualize_areas(screenshot, ore_locations, magma_locations, x, y, width, height)

    # 如果检测到角色位置
    if role_locations:
        # 计算角色中心位置（像素坐标）
        x1, y1, x2, y2 = role_locations[0]
        center_x = int((x1 + x2) / 2)
        center_y = int((y1 + y2) / 2)

        # 创建安全区域网格
        grid = create_grid_from_mask(safe_mask)

        # 将角色位置转换为网格坐标
        current_grid_pos = pixel_to_grid((center_x, center_y), width, height)

        # 如果还没有目标边缘或路径，寻找新的边缘并规划路径
        if not hasattr(move_in_safe_area, 'target_edge') or not hasattr(move_in_safe_area, 'current_path') or len(
                move_in_safe_area.current_path) == 0:
            print("寻找新的安全区域边缘...")
            move_in_safe_area.target_edge = find_safe_edge(grid, current_grid_pos)
            print(f"找到边缘位置: 网格({move_in_safe_area.target_edge[0]}, {move_in_safe_area.target_edge[1]})")

            # 转换为像素坐标
            target_pixel_pos = grid_to_pixel(move_in_safe_area.target_edge)
            print(f"边缘位置（像素）: ({target_pixel_pos[0]}, {target_pixel_pos[1]})")

            print("规划路径...")
            move_in_safe_area.current_path = a_star_search(grid, current_grid_pos, move_in_safe_area.target_edge)
            print(f"路径长度: {len(move_in_safe_area.current_path)} 步")

            # 绘制目标边缘
            cv2.circle(marked_image, target_pixel_pos, 10, (255, 0, 0), -1)  # 蓝色圆点标记目标

            # 如果找到了路径，绘制路径
            if move_in_safe_area.current_path:
                pixel_path = [grid_to_pixel(pos) for pos in move_in_safe_area.current_path]
                for i in range(len(pixel_path) - 1):
                    cv2.line(marked_image, pixel_path[i], pixel_path[i + 1], (255, 0, 0), 2)

        # 如果有路径，执行下一步移动
        if hasattr(move_in_safe_area, 'current_path') and len(move_in_safe_area.current_path) > 0:
            next_grid_point = move_in_safe_area.current_path[0]
            next_pixel_point = grid_to_pixel(next_grid_point)

            # 计算需要移动的方向
            dx = next_pixel_point[0] - center_x
            dy = next_pixel_point[1] - center_y

            # 确定移动方向
            if abs(dx) > abs(dy):
                direction = (1, 0) if dx > 0 else (-1, 0)
            else:
                direction = (0, 1) if dy > 0 else (0, -1)

            # 获取对应键盘操作并执行
            key_map = {
                (0, -1): 'up',
                (1, 0): 'right',
                (0, 1): 'down',
                (-1, 0): 'left'
            }

            key = key_map.get(direction)
            if key:
                pydirectinput.keyDown(key)
                time.sleep(0.1)
                pydirectinput.keyUp(key)
                print(f"向 {['上', '右', '下', '左'][DIRECTIONS.index(direction)]} 移动")

            # 检查是否到达下一个点
            if abs(center_x - next_pixel_point[0]) < STEP_SIZE // 2 and abs(
                    center_y - next_pixel_point[1]) < STEP_SIZE // 2:
                # 移除已走过的路径点
                move_in_safe_area.current_path = move_in_safe_area.current_path[1:]

            # 绘制当前位置和下一个目标点
            cv2.circle(marked_image, (center_x, center_y), 8, (0, 0, 255), -1)  # 红色圆点标记当前位置
            if move_in_safe_area.current_path:
                next_point = grid_to_pixel(move_in_safe_area.current_path[0])
                cv2.circle(marked_image, next_point, 8, (0, 255, 0), -1)  # 绿色圆点标记下一个目标

        # 绘制原始检测框
        for r in results:
            im_array = r.plot()
            combined_image = cv2.addWeighted(im_array, 0.7, marked_image, 0.7, 0)
            cv2.imshow('YOLOv8 Inference', combined_image)
            cv2.waitKey(1)
    return


def main():
    x = 0
    y = 87
    width = 842
    height = 570

    # 加载模型
    model = YOLO('best.pt')

    while True:
        move_in_safe_area(x, y, width, height, model)

        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

    cv2.destroyAllWindows()


if __name__ == "__main__":
    main()